# 导入必要的库和模块
import gradio as gr
import joblib
from PIL import Image
import numpy as np

# 尝试加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
try:
    knn_model = joblib.load('best_knn_model.pkl')
except Exception as e:
    print(f"Error loading model: {e}")

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(drawing):
    try:
        image = Image.fromarray(drawing)
        image = image.convert("L")
        image = image.resize((8, 8), Image.Resampling.LANCZOS)
        image_array = np.array(image).flatten()
        image_array = image_array.reshape(1, -1)
        prediction = knn_model.predict(image_array)
        # 进行预测
        return int(prediction[0])
    except Exception as e:
        return f"Error during prediction: {e}"

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="绘制图像"),
    outputs="label",
    title="KNN 手写数字识别",
    description="在手写板上绘制一个数字并获取预测结果。"
)
# 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch()
